﻿#include "infer_cls_v2_o.h"
#include "common/define.h"
#include <assert.h>
#include <fstream>
#include <vector>
#include <string>

using namespace ov;

namespace txr_algo_dlm_cls
{ 

	/**
	 * @brief 从文件中读取数据
	 *
	 * @param file_path 文件的路径。
	 * @return 返回包含文件内容的std::vector<char>。
	 */
	std::vector<char> get_file_data(const std::string& file_path) {
		std::ifstream file(file_path, std::ios_base::binary);
		assert(file.good());
		file.seekg(0, std::ios_base::end);
		int file_len = file.tellg();
		file.seekg(0, std::ios_base::beg);
		std::vector<char> data(file_len, 0);
		file.read(data.data(), file_len);
		file.close();
		return data;
	}

    /**
     * @brief 对数据执行简单解密操作
     * @param data 需要解密的数据
     */
    void decrypt_data(std::vector<char>& data) {
        char key = 0x1;
        for (size_t i = 0; i < data.size(); ++i) {
            data[i] ^= key;
        }
    }
	
	    /**
     * @brief 解密文件并提取XML和BIN数据
     * @param merged_path 加密文件路径
     * @param info 存储解密信息的结构体
     * @param xml_data 存储XML数据的向量
     * @param bin_data 存储BIN数据的向量
     * @return 解密成功返回true,失败返回false
     */
    bool decryptFileTo_XB_buffer(
        const std::string& merged_path,
        st_encrypt_info_cls2_o& info, 
        std::vector<char>& xml_data,
        std::vector<char>& bin_data)
    {
        // 检查文件后缀名
        if (merged_path.size() < 6 || merged_path.substr(merged_path.size() - 6) != DATS_EXTENSION) {
            throw std::runtime_error("The provided file is not a .ocls2 file.");
        }

        // 读取合并的数据文件
        std::vector<char> merged_data = get_file_data(merged_path);

        // 检查数据大小是否足够包含头部信息
        if (merged_data.size() < sizeof(st_encrypt_header)) {
            throw std::runtime_error("Merged data is too small to contain the header.");
        }

        // 读取头部信息
        st_encrypt_header header;
        std::memcpy(&header, merged_data.data(), sizeof(st_encrypt_header));

        // 检查数据大小是否足够包含所有部分
        if (merged_data.size() < sizeof(st_encrypt_header) + header.struct_size + header.xml_size + header.bin_size) {
            throw std::runtime_error("Merged data is too small to contain all parts.");
        }

        // 复制结构体信息
        std::memcpy(&info, merged_data.data() + sizeof(st_encrypt_header), header.struct_size);

        // 提取并解密XML数据
        xml_data.assign(
            merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size,
            merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size
        );
        decrypt_data(xml_data);

        // 提取并解密BIN数据
        bin_data.assign(
            merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size,
            merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size + header.bin_size
        );
        decrypt_data(bin_data);

        return true;
    }

	    /**
     * @brief 将权重数据从std::vector<char>转换为OpenVINO的ov::Tensor对象
     *
     * @param weightsBuffer 包含权重数据的字符向量
     * @return 返回一个数据类型为u8且大小与权重数据匹配的ov::Tensor对象
     */
    ov::Tensor vector_to_tensor(const std::vector<char>& weightsBuffer) {
        size_t weightsSize = weightsBuffer.size();
        ov::Tensor weights(ov::element::u8, {weightsSize});
        std::memcpy(weights.data(), weightsBuffer.data(), weightsSize);
        return weights;
    }

	/**
	 * @brief 分类推理类的构造函数
	 */
	InferClsV2::InferClsV2()
	{
	}

	/**
	 * @brief 分类推理类的析构函数
	 */
	InferClsV2::~InferClsV2()
	{
	}

	/**
	 * @brief 初始化推理环境
	 * @param gpu_id GPU设备ID
	 * @param dats_path 模型数据路径
	 * @return 初始化是否成功
	 */
	bool InferClsV2::Init(
		int gpu_id, 
		const char* dats_path)
	{
		// 初始化OpenVINO
		core = Core();

		// 解密并编译模型
		const std::string file_path = dats_path;
		std::vector<char> xml_buffer = {};
		std::vector<char> bin_buffer = {};

		decryptFileTo_XB_buffer(
			dats_path,
			m_info,
			xml_buffer,
			bin_buffer);

		std::string model_xml(xml_buffer.begin(), xml_buffer.end());
		ov::Tensor bin_tensor = vector_to_tensor(bin_buffer);

		try {
			std::shared_ptr<ov::Model> model = core.read_model(model_xml, bin_tensor);
			compiled_model = core.compile_model(model, "AUTO");
		}
		catch (const std::exception& e) {
			throw;
		}

		// 创建推理请求
		size_t num_requests = m_info.batch_size;
		infer_requests.resize(num_requests);
		for (size_t i = 0; i < num_requests; ++i) {
			infer_requests[i] = compiled_model.create_infer_request();
		}

		return true;
	}

	/**
	 * @brief 加载TensorRT推理引擎
	 * @param v_engine_data 引擎数据
	 * @return 加载是否成功
	 */
	bool InferClsV2::LoadEngine(std::vector<char> v_engine_data)
	{

		return true;
	}

	/**
	 * @brief 图像预处理函数,将输入图像转换为模型所需格式
	 * @param p_unit 检测单元数组指针,包含输入图像数据
	 * @param num 图像数量
	 * @param target_w 目标宽度
	 * @param target_h 目标高度  
	 * @param mean 归一化均值数组,长度为3对应RGB三通道
	 * @param std 归一化标准差数组,长度为3对应RGB三通道
	 * @param processed_data 输出处理后的图像数据
	 */
	void PreprocessToMat(const st_detect_unit* p_unit, int num, 
						int target_w, int target_h, 
						const float mean[3], const float std[3],
						std::vector<std::vector<float>>& processed_data)
	{
		processed_data.clear();
		processed_data.reserve(num);

		for (int i = 0; i < num; ++i)
		{
			const st_img_rgb& img = p_unit[i].img;

			int src_w = img.w, src_h = img.h;
			cv::Mat src_image(src_h, src_w, CV_32FC3);

			// 根据输入数据类型进行处理
			for (int c = 0; c < 3; ++c)
			{
				for (int y = 0; y < src_h; ++y)
				{
					for (int x = 0; x < src_w; ++x)
					{
						float pixel_value = 0.0f;
						if (img.pp_rgbf[c]) {
							// float类型数据直接使用
							pixel_value = img.pp_rgbf[c][y * src_w + x];
						}
						else if (img.pp_rgb8[c]) {
							// uint8类型数据转换为float
							pixel_value = static_cast<float>(img.pp_rgb8[c][y * src_w + x]);
						}
						else if (img.pp_rgb16[c]) {
							// uint16类型数据转换为float
							pixel_value = static_cast<float>(img.pp_rgb16[c][y * src_w + x]);
						}
						else {
							std::cerr << "Warning: No valid data found for channel " << c << " of image " << i << std::endl;
							continue;
						}
						src_image.at<cv::Vec3f>(y, x)[c] = pixel_value;
					}
				}
			}

			// 缩放图像到目标尺寸
			cv::Mat resized_image;
			cv::resize(src_image, resized_image, cv::Size(target_w, target_h), 0, 0, cv::INTER_LINEAR);

			// 标准化
			resized_image.forEach<cv::Vec3f>([&](cv::Vec3f& pixel, const int* position) {
				pixel[0] = (pixel[0] - mean[0]) / std[0];  // R通道
				pixel[1] = (pixel[1] - mean[1]) / std[1];  // G通道
				pixel[2] = (pixel[2] - mean[2]) / std[2];  // B通道
			});

			// 转换为 OpenVINO 输入格式：NCHW -> std::vector<float>
			std::vector<float> nchw_data(3 * target_h * target_w);
			for (int c = 0; c < 3; ++c) {
				for (int h = 0; h < target_h; ++h) {
					for (int w = 0; w < target_w; ++w) {
						nchw_data[c * target_h * target_w + h * target_w + w] = resized_image.at<cv::Vec3f>(h, w)[c];
					}
				}
			}

			processed_data.push_back(nchw_data);
		}
	}


	/**
	 * @brief 对单张图像进行预处理
	 * @param img 输入图像
	 * @param host_buffer 主机端缓存(Mat格式)
	 * @param net_w 目标宽度
	 * @param net_h 目标高度
	 * @param mean 均值数组
	 * @param std 标准差数组
	 */
	void PreprocessSingleImage(const st_img_rgb& img, 
							cv::Mat& host_buffer,
							int net_w, 
							int net_h,
							const float* mean,
							const float* std) {
		// 输入验证
		if (!img.pp_rgb8[0] && !img.pp_rgb16[0] && !img.pp_rgbf[0]) {
			throw std::runtime_error("无效的输入图像数据");
		}
		if (net_w <= 0 || net_h <= 0 || img.w <= 0 || img.h <= 0) {
			throw std::runtime_error("无效的图像尺寸");
		}

		cv::Mat input_mat;
		// 创建输入Mat
		if (img.pp_rgb8[0]) {
			std::vector<cv::Mat> channels;
			for(int c = 0; c < 3; c++) {
				if (!img.pp_rgb8[c]) throw std::runtime_error("缺少RGB通道数据");
				channels.push_back(cv::Mat(img.h, img.w, CV_8UC1, img.pp_rgb8[c]));
			}
			cv::merge(channels, input_mat);
		}
		else if (img.pp_rgb16[0]) {
			std::vector<cv::Mat> channels;
			for(int c = 0; c < 3; c++) {
				if (!img.pp_rgb16[c]) throw std::runtime_error("缺少RGB通道数据");
				channels.push_back(cv::Mat(img.h, img.w, CV_16UC1, img.pp_rgb16[c]));
			}
			cv::merge(channels, input_mat);
			// 添加值范围限制
			cv::Mat temp;
			input_mat.convertTo(temp, CV_8UC3, 1.0/256.0);
			input_mat = temp;
		}
		else if (img.pp_rgbf[0]) {
			std::vector<cv::Mat> channels;
			for(int c = 0; c < 3; c++) {
				if (!img.pp_rgbf[c]) throw std::runtime_error("缺少RGB通道数据");
				channels.push_back(cv::Mat(img.h, img.w, CV_32FC1, img.pp_rgbf[c]));
			}
			cv::merge(channels, input_mat);
			// 添加值范围限制
			cv::Mat temp;
			input_mat = cv::max(cv::min(input_mat, 1.0f), 0.0f);
			input_mat.convertTo(temp, CV_8UC3, 255.0);
			input_mat = temp;
		}

		// 计算resize尺寸
		float ratio = float(img.w) / float(img.h);
		int resize_w = std::min(net_w, int(std::ceil(net_h * ratio)));
		int resize_h = net_h;

		// 一步到位完成resize和padding
		cv::Mat padded = cv::Mat::zeros(net_h, net_w, CV_8UC3);
		cv::Mat roi = padded(cv::Rect((net_w - resize_w)/2, (net_h - resize_h)/2, resize_w, resize_h));
		cv::resize(input_mat, roi, cv::Size(resize_w, resize_h));

		// 转换为float并归一化
		padded.convertTo(host_buffer, CV_32FC3, 1.0/255.0);
		std::vector<cv::Mat> channels(3);
		cv::split(host_buffer, channels);
		for(int c = 0; c < 3; c++) {
			channels[c] = (channels[c] - mean[c]) / std[c];
		}
		cv::merge(channels, host_buffer);
	}

	void postProcess(st_detect_unit* p_unit, const float* output, const int& out_size)
	{
		// 找到输出中概率最大值及对应类别索引
		auto max_pos = std::max_element(output, output + out_size);

		// 将结果存储到结构体中
		p_unit->result.cls = static_cast<tk_int32>(std::distance(output, max_pos));  // 最大值的索引
		p_unit->result.prob = *max_pos;  // 最大值的概率
	}

	/**
	 * @brief 执行图像分类检测
	 * @param p_unit 检测单元数组
	 * @param num 检测单元数量
	 */
	void InferClsV2::Detect(st_detect_unit* p_unit, int num)
	{
		// 检查输入数量是否超过batch size
		int batch_size = m_info.batch_size;
		if (num > batch_size)
		{
			num = batch_size;
		}

		// 获取网络输入尺寸
		int net_w = m_info.image_width,
			net_h = m_info.image_height,
			net_c = m_info.input_channel;

		std::vector<std::vector<float>> processed_data;

		// 图像预处理循环
		float mean[3] = { m_info.img_mean[0], m_info.img_mean[1], m_info.img_mean[2] };
		float std[3] = { m_info.img_std[0], m_info.img_std[1], m_info.img_std[2] };
		
		PreprocessToMat(p_unit, num, net_w, net_h, mean, std, processed_data);

		// 推理
		for (size_t i = 0; i < batch_size; ++i) {
			ov::Tensor input_tensor(ov::element::f32, 
									{1, 3, static_cast<size_t>(net_h), static_cast<size_t>(net_w)}, 
									processed_data[i].data());

			infer_requests[i].set_input_tensor(input_tensor);
			infer_requests[i].infer();

			const ov::Tensor output_tensor = infer_requests[i].get_output_tensor();
			const float* output_data = output_tensor.data<float>();

			// 后处理，将结果存储到 st_detect_unit
			postProcess(&p_unit[i], output_data, 2);  // 假设输出类别数为2
		}
	}



}//txr_algo_dlm_cls